#include "yolov3_clip_detector.h"
#include <algorithm>
#include "acl/acl.h"
#include <dirent.h>
#include <math.h>
#include <cstdlib>
#include <chrono>
#include <unistd.h>

Yolov3CLIPDetector::Yolov3CLIPDetector()
{
  std::cout<<"-- Creat instence of Resource Yolov3CLIPDetector"<<std::endl;
}

void Yolov3CLIPDetector::DestroyResource()
{
    done = true;
    queue_input.notify_all();
    queue_pre.notify_all();
    queue_post.notify_all();

    for (auto& th : threads) {
        if (th.joinable()) {
            th.join();
        }
    }

    if (modelDesc != nullptr) {
        (void)aclmdlDestroyDesc(modelDesc);
        modelDesc = nullptr;
    }

    aclError ret = aclmdlUnload(modelId);
    if (ret != ACL_SUCCESS) {
        ACLLITE_LOG_ERROR("Unload model error:%d", ret);
    }

    
    for (size_t i = 0; i < aclmdlGetDatasetNumBuffers(input); ++i) {
        aclDataBuffer* dataBuffer = aclmdlGetDatasetBuffer(input, i);
        aclDestroyDataBuffer(dataBuffer);
        dataBuffer = nullptr;
    }
    aclmdlDestroyDataset(input);
    input = nullptr;

    for (size_t i = 0; i < aclmdlGetDatasetNumBuffers(output); ++i) {
        aclDataBuffer* dataBuffer = aclmdlGetDatasetBuffer(output, i);
        void* data = aclGetDataBufferAddr(dataBuffer);
        (void)aclrtFree(data);
        (void)aclDestroyDataBuffer(dataBuffer);
        dataBuffer = nullptr;
    }

    (void)aclmdlDestroyDataset(output);
    output = nullptr;
    if (outputHostData0 != nullptr){
        aclrtFreeHost(outputHostData0);
        outputHostData0 = nullptr;
        std::cout<<"destroy outputHostData0 success"<<std::endl;
    }
    if (outputHostData1 != nullptr){
        aclrtFreeHost(outputHostData1);
        outputHostData1 = nullptr;
        std::cout<<"destroy outputHostData1 success"<<std::endl;
    }

    aclrtFree(modelInput0Buf);
    aclrtFree(modelInput1Buf);

    ret = aclrtDestroyContext(context);
    if (ret != ACL_SUCCESS) {
        ACLLITE_LOG_ERROR("destroy context failed, errorCode is : %d", ret);
    }
    context = nullptr;

    ret = aclrtResetDevice(deviceId);
    if (ret != ACL_SUCCESS) {
        ACLLITE_LOG_ERROR("reset device failed, errorCode is : %d", ret);
    }
    ACLLITE_LOG_INFO("Reset device %d ok", deviceId); 

    ret = aclFinalize();
    if (ret != ACL_SUCCESS) {
        ACLLITE_LOG_ERROR("finalize acl failed, errorCode is : %d", ret);
    }
    ACLLITE_LOG_INFO("Finalize acl ok");
}

void Yolov3CLIPDetector::wait_for_completion() {
    std::unique_lock<std::mutex> lock(results_mutex);
    conda.wait(lock, [this]{ return this->task_counter == 0; });
}

Yolov3CLIPDetector::~Yolov3CLIPDetector()
{
  DestroyResource();
}

AclLiteError Yolov3CLIPDetector::CreateInput()
{
    aclError aclRet = aclrtMalloc(&modelInput0Buf, (size_t)(modelInputSize), ACL_MEM_MALLOC_HUGE_FIRST);
    if (aclRet != ACL_SUCCESS) {
        ACLLITE_LOG_ERROR("malloc device data buffer failed, aclRet is %d", aclRet);
        return ACLLITE_ERROR;
    }

    aclRet = aclrtMalloc(&modelInput1Buf, (size_t)(modelInputSize), ACL_MEM_MALLOC_HUGE_FIRST);
    if (aclRet != ACL_SUCCESS) {
        ACLLITE_LOG_ERROR("malloc device data buffer failed, aclRet is %d", aclRet);
        return ACLLITE_ERROR;
    }

    std::vector<DataInfo> inputData = {{modelInput0Buf, modelInputSize}, {modelInput1Buf, modelInputSize}};

    uint32_t dataNum = aclmdlGetNumInputs(modelDesc);
    if (dataNum == 0) {
        std::cout<<"Create input failed for no input data"<<std::endl;
        return ACLLITE_ERROR_INVALID_ARGS;
    }

    input = aclmdlCreateDataset();
    if (input == nullptr) {
        std::cout<<"Create input failed for create dataset failed"<<std::endl;
        return ACLLITE_ERROR_CREATE_DATASET;
    }

    for (uint32_t i = 0; i < inputData.size(); i++) {
        size_t modelInputSize_1 = aclmdlGetInputSizeByIndex(modelDesc, i);
        std::cout<<"modelInputSize_1: "<<modelInputSize_1<<"inputData[i].size: "<<inputData[i].size<<std::endl;
        if (modelInputSize != inputData[i].size) {
            ACLLITE_LOG_WARNING("Input size verify failed "
                                "input[%d] size: %ld, provide size : %d",
                                i, modelInputSize, inputData[i].size);
        }
        aclDataBuffer* dataBuf = aclCreateDataBuffer(inputData[i].data, inputData[i].size);
        if (inputData[i].data == nullptr) {
            ACLLITE_LOG_ERROR("Create data buffer error");
            return ACLLITE_ERROR_CREATE_DATA_BUFFER;
        }

        aclError ret = aclmdlAddDatasetBuffer(input, dataBuf);
        if (ret != ACL_SUCCESS) {
            ACLLITE_LOG_ERROR("Add dataset buffer error %d", ret);
            ret = aclDestroyDataBuffer(dataBuf);
            if (ret != ACL_SUCCESS) {
                ACLLITE_LOG_ERROR("Destroy dataset buffer error %d", ret);
            }
            dataBuf = nullptr;
            return ACLLITE_ERROR_ADD_DATASET_BUFFER;
        }
    }
    return ACLLITE_OK;
}

void Yolov3CLIPDetector::initThread(){
    for (u_int32_t i=0; i<preprocess_thread_num;i++){
        threads.emplace_back(&Yolov3CLIPDetector::Preprocess, this);
    }
    for (u_int32_t i=0; i<preprocess_thread_num;i++){
        threads.emplace_back(&Yolov3CLIPDetector::modelExecute, this);
    }
    for (u_int32_t i=0; i<postprocess_thread_num;i++){
        threads.emplace_back(&Yolov3CLIPDetector::Postprocess, this);
    }
}

bool Yolov3CLIPDetector::init(const char * ModelPath)
{
    aclError ret = aclInit("");
    if (ret != ACL_SUCCESS) {
        std::cout<<"Acl init failed, errorcode is: "<< ret<<std::endl;
        return false;
    }
    std::cout<<"Acl init ok"<<std::endl;

    // open device
    ret = aclrtSetDevice(deviceId);
    if (ret != ACL_SUCCESS) {
        std::cout<<"Acl open device "<<deviceId<<" failed, errorCode is : "<<ret<<std::endl;
        return false;
    }
    std::cout<<"Open device "<<deviceId<<" ok"<<std::endl;

    ret = aclrtCreateContext(&context, deviceId);
    if (ret != ACL_SUCCESS) {
        std::cout<<"Create acl context failed, errorCode is : "<<ret<<std::endl;
        return false;
    }

    initThread();
    
    // Gets whether the current application is running on host or Device
    ret = aclrtGetRunMode(& g_runMode_);
    if (ret != ACL_SUCCESS) {
        std::cout<<"acl get run mode failed"<<std::endl;
        return false;
    }

    ret = aclmdlLoadFromFile(ModelPath, &modelId);
    if (ret != ACL_SUCCESS) {
        ACLLITE_LOG_ERROR("Load model(%s) from file return %d",
                          ModelPath, ret);
        return ACLLITE_ERROR_LOAD_MODEL;
    }

    modelDesc = aclmdlCreateDesc();
    if (modelDesc == nullptr) {
        std::cout<<"create model description failed"<<std::endl;
        return false;
    }

    ret = aclmdlGetDesc(modelDesc, modelId);
    if (ret != ACL_SUCCESS) {
        std::cout<<"get model description failed, modelId is "<<modelId<< ", errorCode is "<<static_cast<int32_t>(ret)<<std::endl;
        return false;
    }

    std::cout<<"create model description success"<<std::endl;
    // cteate output
    output = aclmdlCreateDataset();
    if (output == nullptr) {
        std::cout<<"can't create dataset, create output false"<<std::endl;
        return false;
    }

    // size_t outputSize = aclmdlGetNumOutputs(modelDesc);
    // for (size_t i = 0; i < outputSize; ++i) {
    const char* name = aclmdlGetOutputNameByIndex(modelDesc, 0);
    std::cout<<"output op name: "<<name<<std::endl;
    buffer_size = aclmdlGetOutputSizeByIndex(modelDesc, 0);
    
    ret = aclrtMalloc(&outputDeviceData0, buffer_size, ACL_MEM_MALLOC_NORMAL_ONLY);
    if (ret != ACL_SUCCESS) {
        std::cout<<"can't malloc buffer, size is"<<buffer_size<<", create output failed"<<std::endl;
        return false;
    }

    aclDataBuffer* outputData = aclCreateDataBuffer(outputDeviceData0, buffer_size);
    if (ret != ACL_SUCCESS) {
        std::cout<<"can't create data buffer, create output failed"<<std::endl;
        aclrtFree(outputDeviceData0);
        return false;
    }

    ret = aclmdlAddDatasetBuffer(output, outputData);
    if (ret != ACL_SUCCESS) {
        std::cout<<"can't add data buffer, create output failed"<<std::endl;
        aclrtFree(outputDeviceData0);
        aclDestroyDataBuffer(outputData);
        return false;
    }

    name = aclmdlGetOutputNameByIndex(modelDesc, 1);
    std::cout<<"output op name: "<<name<<std::endl;
    
    ret = aclrtMalloc(&outputDeviceData1, buffer_size, ACL_MEM_MALLOC_NORMAL_ONLY);
    if (ret != ACL_SUCCESS) {
        std::cout<<"can't malloc buffer, size is"<<buffer_size<<", create output failed"<<std::endl;
        return false;
    }

    outputData = aclCreateDataBuffer(outputDeviceData1, buffer_size);
    if (ret != ACL_SUCCESS) {
        std::cout<<"can't create data buffer, create output failed"<<std::endl;
        aclrtFree(outputDeviceData1);
        return false;
    }

    ret = aclmdlAddDatasetBuffer(output, outputData);
    if (ret != ACL_SUCCESS) {
        std::cout<<"can't add data buffer, create output failed"<<std::endl;
        aclrtFree(outputDeviceData1);
        aclDestroyDataBuffer(outputData);
        return false;
    }
   
    std::cout<<"create model output success"<<std::endl;

    ret = aclrtMallocHost(&outputHostData0, buffer_size);
    if (ret != ACLLITE_OK) {
        std::cout<<"Create model input failed"<<std::endl;
        return false;
    }

    ret = aclrtMallocHost(&outputHostData1, buffer_size);
    if (ret != ACLLITE_OK) {
        std::cout<<"Create model input failed"<<std::endl;
        return false;
    }

    modelInputSize = aclmdlGetInputSizeByIndex(modelDesc, 0);
    std::cout<<"modelInputSize: "<<modelInputSize<<std::endl;
    ret = CreateInput();
    if (ret != ACLLITE_OK) {
        std::cout<<"Create model input failed"<<std::endl;
        return false;
    }
    std::cout<<"create model input success"<<std::endl;
    return true;
}

void* Yolov3CLIPDetector::GetInferenceOutputItem(uint32_t& itemDataSize,aclmdlDataset* inferenceOutput,uint32_t idx)
{
    aclDataBuffer* dataBuffer = aclmdlGetDatasetBuffer(inferenceOutput, idx);
    if (dataBuffer == nullptr) {
        std::cout<<"Get the " <<idx<<" th dataset buffer from model inference output failed"<<std::endl;
        return nullptr;
    }

    void* dataBufferDev = aclGetDataBufferAddr(dataBuffer);
    if (dataBufferDev == nullptr) {
        std::cout<<"Get the "<<idx<<" th dataset buffer address from model inference output failed"<<std::endl;
        return nullptr;
    }

    size_t bufferSize = aclGetDataBufferSize(dataBuffer);
    if (bufferSize == 0) {
        std::cout<<"The "<<idx<<" th dataset buffer size of model inference output is 0"<<std::endl;
        return nullptr;
    }

    void* data = nullptr;
    if (g_runMode_ == ACL_HOST) {
        data = CopyDataDeviceToLocal(dataBufferDev, bufferSize);
        if (data == nullptr) {
            std::cout<<"Copy inference output to host failed"<<std::endl;
            return nullptr;
        }
    } else {
        data = dataBufferDev;
    }
    itemDataSize = bufferSize;
    return data;
}

void* Yolov3CLIPDetector::CopyDataDeviceToLocal(void* deviceData, uint32_t dataSize)
{
    uint8_t* buffer = new uint8_t[dataSize];
    if (buffer == nullptr) {
        std::cout<<"New malloc memory failed"<<std::endl;
        return nullptr;
    }

    aclError aclRet = aclrtMemcpy(buffer, dataSize, deviceData, dataSize, ACL_MEMCPY_DEVICE_TO_HOST);
    if (aclRet != ACL_SUCCESS) {
        std::cout<<"Copy device data to local failed, aclRet is "<< aclRet<<std::endl;
        delete[](buffer);
        return nullptr;
    }

    return (void*)buffer;
}

void Yolov3CLIPDetector::Preprocess()
{
    message_pre inputmsg;
    aclrtSetCurrentContext(context);

    while(!done){
        if (queue_input.try_pop(inputmsg)){
            cv::Mat reiszeMat_visible;
            cv::resize(inputmsg.frame_visible, reiszeMat_visible, cv::Size(g_modelWidth_, g_modelHeight_));
            inputmsg.frame_visible = reiszeMat_visible;

            cv::Mat reiszeMat_infrared;
            cv::resize(inputmsg.frame_infrared, reiszeMat_infrared, cv::Size(g_modelWidth_, g_modelHeight_));
            inputmsg.frame_infrared = reiszeMat_infrared;

            queue_pre.push(inputmsg); 
        }else{
            std::this_thread::sleep_for(std::chrono::milliseconds(10));
        }
    } 
}

AclLiteError Yolov3CLIPDetector::modelExecute()
{
    aclrtSetCurrentContext(context);
    message_pre premsg;
    message_result excutemsg;

    uint32_t dataSize = 0;
    float* outFloatData0;
    float* outFloatData1;
    const uint32_t visible_DataBufId = 0;
    const uint32_t infrared_DataBufId = 1;

    while(!done){
        if (queue_pre.try_pop(premsg)){
            excutemsg.finalOutput.clear();
            cv::Mat reiszeMat_infrared = premsg.frame_infrared;
            cv::Mat reiszeMat_visible = premsg.frame_visible;
            {
                std::lock_guard<std::mutex> lock(output_mutex);
                aclError ret = aclrtMemcpy(modelInput0Buf, modelInputSize, reiszeMat_visible.ptr<uint8_t>(), modelInputSize, ACL_MEMCPY_DEVICE_TO_DEVICE);
                if (ret != ACL_SUCCESS) {
                    ACLLITE_LOG_ERROR("Copy data to device failed, aclRet is %d", ret);
                    break;
                }
                ret = aclrtMemcpy(modelInput1Buf, modelInputSize, reiszeMat_infrared.ptr<uint8_t>(), modelInputSize, ACL_MEMCPY_DEVICE_TO_DEVICE);
                if (ret != ACL_SUCCESS) {
                    ACLLITE_LOG_ERROR("Copy data to device failed, aclRet is %d", ret);
                    break;
                }
               
                // Perform reasoning
                ret = aclmdlExecute(modelId, input, output);
                if (ret != ACL_SUCCESS) {
                    std::cout<<"execute aclmdlExecute false "<<static_cast<int32_t>(ret)<<std::endl;
                    return false;
                }
                
                if (g_runMode_ == ACL_HOST) {
                    ret = aclrtMemcpy(outputHostData0, buffer_size, outputDeviceData0, buffer_size, ACL_MEMCPY_DEVICE_TO_HOST);
                    outFloatData0 = reinterpret_cast < float * > (outputHostData0);
                    ret = aclrtMemcpy(outputHostData1, buffer_size, outputDeviceData1, buffer_size, ACL_MEMCPY_DEVICE_TO_HOST);
                    outFloatData0 = reinterpret_cast < float * > (outputHostData0);
                }
                else
                {
                    outFloatData0 = reinterpret_cast < float * > (outputDeviceData0);
                    outFloatData1 = reinterpret_cast < float * > (outputDeviceData1);
                }       
            }
            std::vector<float> vec_0(outFloatData0, outFloatData0 + 6);
            std::vector<float> vec_1(outFloatData1, outFloatData1 + 6);
            excutemsg.finalOutput.push_back(vec_0);
            excutemsg.finalOutput.push_back(vec_1);
            excutemsg.filename_infrared = premsg.filename_infrared;
            excutemsg.filename_visible = premsg.filename_visible;
            queue_post.push(excutemsg);
            uint32_t dataSize = 0;
        }else{
            std::this_thread::sleep_for(std::chrono::milliseconds(10));
        }
    }
}

void Yolov3CLIPDetector::Postprocess()
{
    float* detectData = nullptr;
    message_result postmsg;
    aclrtSetCurrentContext(context);

    while(!done){
        message_result postresult;
        if (queue_post.try_pop(postmsg)) {
            postresult.filename_infrared = postmsg.filename_infrared;
            postresult.filename_visible = postmsg.filename_visible;
            postresult.finalOutput = postmsg.finalOutput;
            std::lock_guard<std::mutex> lock(results_mutex);
            results.push_back(postresult);
            if (--task_counter == 0) {
                conda.notify_all();
            }
        }else{
            std::this_thread::sleep_for(std::chrono::milliseconds(10));
        }
    }

}

std::vector<message_result> Yolov3CLIPDetector::inference(std::vector<std::string> &imageLists){
    {
        std::lock_guard<std::mutex> lock(results_mutex);
        results.clear();
        task_counter = 0;
    }

    message_pre premsg;
    for (int i=0; i<imageLists.size(); i++){
        cv::Mat img1 = cv::imread(imageLists[i]);
        premsg.filename_visible = imageLists[i];
        premsg.frame_visible = img1;
        premsg.filename_infrared = imageLists[i];
        premsg.frame_infrared = img1;
        queue_input.push(premsg);
        task_counter++;
    }

    wait_for_completion();

    std::vector<message_result> local_results;
    {
        std::lock_guard<std::mutex> lock(results_mutex);
        local_results = results;
    }

    return local_results;

}

